import numpy as np
import platgo as pg


class DTLZ7(pg.Problem):

    def __init__(self, M: int = 3) -> None:
        self.name = 'DTLZ7'
        self.type['multi'], self.type['many'], self.type['real'], self.type['large'], self.type['expensive'] = [True] * 5
        self.M = M
        self.D = M + 19
        lb = [0] * self.D
        ub = [1] * self.D
        self.borders = np.array([lb, ub])
        super().__init__()

    def cal_obj(self, pop: pg.Population) -> None:
        decs = pop.decs
        pop.cv = np.zeros((pop.N, self.D))
        g = 1 + 9 * np.sum(decs[:, (self.M - 1):], axis=1, keepdims=True) / 20
        pop.objv = np.ones((pop.N, self.M))
        pop.objv[:, 0:(self.M - 1)] = decs[:, 0:(self.M - 1)]
        f = (1 + g) * (self.M - np.sum(pop.objv[:, :(self.M - 1)] * (1 + np.sin(3 * np.pi * pop.objv[:, :(self.M - 1)])) / (1 + g), axis=1, keepdims=True))
        pop.objv[:, (self.M - 1):] = f[:, 0:]

    def get_optimal(self) -> np.ndarray:
        # 目标空间均匀采样函数还未完成
        raise NotImplementedError("get optimal has not been implemented")


if __name__ == '__main__':
    d = DTLZ7()
    # pop = pg.Population(decs=np.array([[2, 2.0, 2], [1, 2, 2], [1, 3, 2]]))
    pop = pg.Population(decs=np.random.uniform(0, 1, (10, 7)))
    print(d.borders)
    print(pop.decs)
    d.cal_obj(pop)  # 计算目标函数值
    print(pop.objv)

